﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using Kokomo.Mirror.Ia32;
using Kokomo.PeInspector.Reflection;
using Kokomo.Mirror.PeFormat;
using Kokomo.Mirror;
using Kokomo.PeInspector.Native;

namespace Kokomo.PeInspector
{
	public class DisassemblyDocumentBuilder
	{
		public static void BuildDisassemblyDocument(
			NativeProcedureInfo procedure,
			ICodeWriter writer
			)
		{
			DisassemblyDocumentBuilder docBuilder = new DisassemblyDocumentBuilder(new PeModuleDisassemblyContext(procedure.Module));

			Ia32Disassembly disassembly = procedure.Disassemble();
			if (disassembly != null)
			{
				docBuilder.BuildDisassemblyDocument(
					disassembly,
					writer,
					DisassemblyOptions.IncludeOffset | DisassemblyOptions.IncludeBytes | DisassemblyOptions.LookUpNames
					);
			}
		}
		public static void BuildSourceDocument(
			NativeProcedureInfo procedure,
			ICodeWriter writer
			)
		{
			if (procedure == null) throw new ArgumentNullException("procedure");
			if (writer == null) throw new ArgumentNullException("writer");

			NativePseudocodeDocumentBuilder builder = new NativePseudocodeDocumentBuilder(
				writer
				);

			builder.WriteProcedure(procedure);
			//builder.o(procedure);
		}

		public DisassemblyDocumentBuilder(
			IDisassemblyContext context
			)
		{
			if (context == null) throw new ArgumentNullException("context");
			this.Context = context;
		}

		public IDisassemblyContext Context { get; private set; }

		public event EventHandler<InstructionEventArgs> WritingInstruction;
		protected virtual void OnWritingInstruction(Ia32Instruction instruction)
		{
			if (this.WritingInstruction != null)
				this.WritingInstruction(this, new InstructionEventArgs(instruction));
		}

		public DisassemblyDocumentMetrics BuildDisassemblyDocument(
			Ia32DisassemblyBase disassembly,
			ICodeWriter writer,
			DisassemblyOptions options = DisassemblyOptions.None
			)
		{
			DisassemblyDocumentMetrics metrics = this.CalculateDocumentMetrics(disassembly.Instructions);
			this.BuildDisassemblyDocument(disassembly.Instructions, writer, options, metrics);

			return metrics;
		}

		public DisassemblyDocumentMetrics BuildDisassemblyDocument(
			Ia32Disassembly disassembly,
			ICodeWriter writer,
			DisassemblyOptions options = DisassemblyOptions.None
			)
		{
			DisassemblyDocumentMetrics metrics = this.CalculateDocumentMetrics(disassembly.Instructions);

			foreach (var block in disassembly.Blocks)
			{
				this.BuildDisassemblyDocument(block.Instructions, writer, options, metrics);
				writer.WriteLineBreak();
				writer.WriteLineBreak();
			}

			return metrics;
		}

		public DisassemblyDocumentMetrics BuildDisassemblyDocument(
			Ia32DisassemblyBlock disassemblyBlock,
			ICodeWriter writer,
			DisassemblyOptions options = DisassemblyOptions.None
			)
		{
			DisassemblyDocumentMetrics metrics = this.CalculateDocumentMetrics(disassemblyBlock.Instructions);
			this.BuildDisassemblyDocument(disassemblyBlock.Instructions, writer, options, metrics);

			return metrics;
		}

		private DisassemblyDocumentMetrics CalculateDocumentMetrics(
			IEnumerable<Ia32Instruction> instructions
			)
		{
			int maxMnemonicLength = 0;
			int maxInstructionLength = 0;
			foreach (var instruction in instructions)
			{
				string mnemonicName = instruction.OpCode.Mnemonic.Name;
				if (mnemonicName != null) maxMnemonicLength = Math.Max(maxMnemonicLength, mnemonicName.Length);
				maxInstructionLength = Math.Max(maxInstructionLength, instruction.Length);
			}

			return new DisassemblyDocumentMetrics()
			{
				BytesWidth = maxInstructionLength * 3 + 1,
				InstructionWidth = maxMnemonicLength
			};
		}

		private void BuildDisassemblyDocument(
			IEnumerable<Ia32Instruction> instructions,
			ICodeWriter writer,
			DisassemblyOptions options,
			DisassemblyDocumentMetrics metrics
			)
		{
			if (instructions == null) throw new ArgumentNullException("instructions");
			if (writer == null) throw new ArgumentNullException("writer");

			bool includeOffset = options.HasFlag(DisassemblyOptions.IncludeOffset);
			bool includeBytes = options.HasFlag(DisassemblyOptions.IncludeBytes);
			bool lookUpNames = options.HasFlag(DisassemblyOptions.LookUpNames);

			string addressFormat =
				(this.Context.DefaultAddressSize == 64)
				? "X16"
				: "X8"
				;
			string displacementAnnotationFormat = "; [{0:" + addressFormat + "}] = 0x{1:X8}";

			foreach (var instruction in instructions)
			{
				writer.EnsureNewLine();

				this.OnWritingInstruction(instruction);

				if (includeOffset)
				{
					string label = instruction.Offset.ToString(addressFormat) + ": ";
					writer.WriteText(label);
				}

				if (includeBytes)
				{
					int paddingWidth;
					if (instruction.Bytes != null)
					{
						for (int i = 0; i < instruction.Bytes.Length; i++)
						{
							byte b = instruction.Bytes[i];
							writer.WriteText(string.Format("{0:X2} ", b), CodeTextType.Annotation);
						}
						paddingWidth = metrics.BytesWidth - (instruction.Length) * 3;
					}
					else
					{
						paddingWidth = metrics.BytesWidth;
					}

					writer.WriteText(new string(' ', paddingWidth));
				}

				string mnemonicName = instruction.OpCode.Mnemonic.Name;
				writer.WriteText(mnemonicName, CodeTextType.Keyword);

				if (instruction.Operands.Length > 0)
				{
					writer.WriteText(new string(' ', (metrics.InstructionWidth - mnemonicName.Length) + 2));
					long? displacementAnnotation = null;
					for (int i = 0; i < instruction.Operands.Length; i++)
					{
						if (i != 0) writer.WriteText(", ");

						Operand operand = instruction.Operands[i];

						this.WriteOperand(
							writer,
							operand,
							lookUpNames,
							out displacementAnnotation
							);
					}

					if (displacementAnnotation.HasValue && (this.Context != null))
					{
						int int32Value;
						string stringValue;
						try
						{
							if (this.Context.ReadData(displacementAnnotation.Value, out int32Value, out stringValue))
							{
								if (int32Value != 0)
								{
									string comment = string.Format(displacementAnnotationFormat,
										displacementAnnotation.Value,
										int32Value
										);
									if (!string.IsNullOrEmpty(stringValue))
										comment += string.Format(" \"{0}\"", stringValue);

									writer.WriteText("     ");
									writer.WriteText(comment, CodeTextType.Comment);
								}
							}
						}
						catch
						{
							// Possibly an invalid memory location
						}
					}
				}
			}
		}

		private void WriteOperand(
			ICodeWriter writer,
			Operand operand,
			bool lookUpSymbolName,
			out long? displacementAnnotation
			)
		{
			displacementAnnotation = null;
			ImmediateOperand immediate = operand as ImmediateOperand;
			if (immediate != null)
			{
				this.WriteAddressOrSymbol(writer, immediate.Value, lookUpSymbolName);
				displacementAnnotation = (long)immediate.Value;
			}
			else
			{
				IndirectOperand indirect = operand as IndirectOperand;
				if (indirect != null)
				{
					if (indirect.SegmentOverride != null)
					{
						writer.WriteText(indirect.SegmentOverride.Name, CodeTextType.TypeName);
						writer.WriteText(":", CodeTextType.Operator);
					}
					writer.WriteText("[", CodeTextType.Operator);
					this.WriteOperand(
						writer,
						indirect.BaseOperand,
						lookUpSymbolName,
						out displacementAnnotation
						);
					if (indirect.Displacement > 0)
					{
						writer.WriteText("+", CodeTextType.Operator);
						writer.WriteText(string.Format("0{0:X}h", indirect.Displacement));
					}
					else if (indirect.Displacement < 0)
					{
						writer.WriteText("-", CodeTextType.Operator);
						writer.WriteText(string.Format("0{0:X}h", -indirect.Displacement));
					}
					writer.WriteText("]", CodeTextType.Operator);
				}
				else
				{
					IndexedOperand indexed = operand as IndexedOperand;
					if (indexed != null)
					{
						if (indexed.SegmentOverride != null)
						{
							writer.WriteText(indexed.SegmentOverride.Name, CodeTextType.TypeName);
							writer.WriteText(":", CodeTextType.Operator);
						}

						writer.WriteText("[", CodeTextType.Operator);
						this.WriteOperand(
							writer,
							indexed.BaseOperand,
							lookUpSymbolName,
							out displacementAnnotation);
						if (indexed.IndexOperand != null)
						{
							writer.WriteText("+", CodeTextType.Operator);
							long? indexAnnotation;
							this.WriteOperand(
								writer,
								indexed.IndexOperand,
								lookUpSymbolName,
								out indexAnnotation
								);
							if (indexed.Scale != 1)
							{
								writer.WriteText("*", CodeTextType.Operator);
								writer.WriteText(string.Format("0{0:X}h", indexed.Scale));
							}
						}
						if (indexed.Displacement > 0)
						{
							writer.WriteText("+", CodeTextType.Operator);
							writer.WriteText(string.Format("0{0:X}h", indexed.Displacement));
						}
						else if (indexed.Displacement < 0)
						{
							writer.WriteText("-", CodeTextType.Operator);
							writer.WriteText(string.Format("0{0:X}h", -indexed.Displacement));
						}
						writer.WriteText("]", CodeTextType.Operator);
					}
					else
					{
						JumpTarget jumpTarget = operand as JumpTarget;
						if (jumpTarget != null)
						{
							this.WriteAddressOrSymbol(writer, (ulong)jumpTarget.TargetOffset, lookUpSymbolName);
							displacementAnnotation = jumpTarget.TargetOffset;
						}
						else
						{
							Register register = operand as Register;
							if (register != null)
							{
								writer.WriteText(register.Name, CodeTextType.TypeName);
							}
						}
					}
				}
			}
		}

		private void WriteAddressOrSymbol(
			ICodeWriter writer,
			ulong address,
			bool lookUpSymbolName
			)
		{
			SymbolicReference reference = null;

			if (lookUpSymbolName && this.Context != null)
				reference = this.Context.GetSymbolicReference((long)address);

			//if (this.Module.IsValidAddress(displacement))
			//{
			//    uint rva = this.Module.AbsoluteToRva(displacement);
			//    reference = this.Module.GetSymbolicReferenceForRva(rva);
			//}
			if (
				(reference != null)
				&& !string.IsNullOrEmpty(reference.SymbolName)
				)
			{
				string modulePath =
					((reference.ModuleName != null) && (this.Context != null))
					? this.Context.ResolveModule(reference.ModuleName)
					: null
					;
				if (modulePath != null)
				{
					writer.WriteMemberReference(new PeModuleRef() { fileName = modulePath }, reference.ModuleName);
					writer.WriteText("!");
				}

				writer.WriteMemberReference(reference, reference.SymbolName);

				if (reference.Offset < 0)
				{
					writer.WriteText("-", CodeTextType.Operator);
					writer.WriteText("0" + (-reference.Offset).ToString("X"));
				}
				else if (reference.Offset > 0)
				{
					writer.WriteText("+", CodeTextType.Operator);
					writer.WriteText("0" + reference.Offset.ToString("X"));
				}
			}
			else
			{
				//string format = (address <= (1 << 32))
				//    ? "X16"
				//    : "X8"
				//    ;
				string format = "0{0:X}h";
				string text = string.Format(format, address);

				NativeProcedureInfo procedure = this.Context.GetProcedure((long)address);
				if (procedure != null)
					writer.WriteMemberReference(procedure, text);
				else
					writer.WriteText(text);
				//writer.WriteText(string.Format(format, address));
			}

		}
	}
}
